package 算法模板;

import java.util.Arrays;

/**
 * 给定一组 n 人（编号为 1, 2, ..., n）， 我们想把每个人分进任意大小的两组。
 * 每个人都可能不喜欢其他人，那么他们不应该属于同一组。
 * 给定整数 n 和数组 dislikes ，其中 dislikes[i] = [ai, bi] ，表示不允许
 * 将编号为 ai 和  bi的人归入同一组。当可以用这种方法将所有人分进两组时，返回 true；否则返回 false。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/possible-bipartition
 */
public class 染色法_邻接表_判断二分图 {

    public static void main(String[] args) {

        System.out.println(new 染色法_邻接表_判断二分图().possibleBipartition(4, new int[][]{{1, 2}, {1, 3}, {2, 4}}));

    }

    // 存储是某个节点所对应的边的集合（链表）的头结点
    int[] he;
    // 访问某一条边指向的节点
    int[] e;
    // 是以链表的形式进行存边，该数组就是用于找到下一条边
    int[] ne;
    // 用于记录某条边的权重为多少
    int[] w;
    // 用来对边进行编号
    int idx;

    // 邻接表用于边数量比较少的情况
    private void add(int a, int b, int c) {
        e[idx] = b;
        ne[idx] = he[a];
        he[a] = idx;
        w[idx] = c;
        idx++;
    }

    /**
     * 定义 0 表示未染色，1 表示红色，2 表示蓝色，3 - color 是 color 的相对颜色
     */
    public boolean possibleBipartition(int n, int[][] dislikes) {
        // 节点数量
        he = new int[n + 1];
        // 边的数量，这里是双向边，需要 *2
        e = new int[dislikes.length * 2 + 1];
        ne = new int[dislikes.length * 2 + 1];
        w = new int[dislikes.length * 2 + 1];

        Arrays.fill(he, -1);
        int[] colors = new int[n + 1];
        for (int[] dislike : dislikes) {
            add(dislike[0], dislike[1], 0);
            add(dislike[1], dislike[0], 0);
        }

        for (int i = 1; i < colors.length; i++) {
            if (colors[i] == 0 &&
                    !dfs(i, 1, colors)) {
                return false;
            }
        }
        return true;
    }

    // 将 color[u] 的位置染成 color
    // 那么需要将 color[u] 的邻接节点染成 3-color
    private boolean dfs(int u, int color, int[] colors) {
        colors[u] = color;

        for (int i = he[u]; i != -1; i = ne[i]) {
            int j = e[i];
            if (colors[j] == 0 &&
                    !dfs(j, 3 - color, colors)) {
                return false;
            } else if (colors[j] == color) {
                return false;
            }
        }

        return true;
    }

}
